#include "Planif.hh"
#include <vector>
#include <map>
#include <algorithm>
#include <limits>
#include "Bug.h"
#include "munkres.h"
#include <map>
#include "Bug.h"
#include <iostream>


struct LocIndex : public std::map<uLocation,unsigned>
{
  std::vector<uLocation> reverse_index;

  LocIndex() : reverse_index(0) {}
  
  
  void insert(uLocation loc)
  {
    std::map<uLocation,unsigned>::iterator it = find(loc);
    if(it == end())
      {
	std::map<uLocation,unsigned>::insert(std::make_pair(loc,reverse_index.size()));
	reverse_index.push_back(loc);
      }
  }
	
  unsigned get_index(uLocation loc)
  {
    std::map<uLocation,unsigned>::iterator it = find(loc);
    return it->second;
  }

  uLocation get_location(unsigned index)
  {
    return reverse_index.at(index);
  }

  
};


std::set<std::pair<uLocation,uLocation> > 
get_an_optimal_plan(std::vector<Option> & options, bool minimize_cost)
{
  // build matrix
  LocIndex ressources;
  LocIndex tasks;
  
  double infinity = 1.0;
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    {
      ressources.insert(it->from);
      tasks.insert(it->to);
      infinity += fabs(it->cost);
      
    }


  signed nrows = ressources.size();
  signed ncols = tasks.size();
  Matrix<double> matrix(nrows,ncols);

  // Initialize matrix with inifinite costs.
  for (signed row = 0 ; row != nrows ; row++ ) {
    for (signed col = 0 ; col != ncols ; col++ ) {
      matrix(row,col) = infinity;
    }
  }

  // discount provided options  
  // to avoid selecting non provided options
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    {
      unsigned row = ressources.get_index(it->from);
      unsigned col = tasks.get_index(it->to);
      matrix(row,col) = 0;
    }


  // fill the cost matrix
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    {
      signed row = ressources.get_index(it->from);
      signed col = tasks.get_index(it->to);

      bug << "Insert " << it->from << " -> " << it->to << " at cost " << it->cost << std::endl;

      
      if(minimize_cost)
	{
	  matrix(row,col) += it->cost;
	}
      else
	{
	  matrix(row,col) -= it->cost;
	}
    }



  // Display solved matrix.
  bug << "cost matrix" << std::endl;
  for ( int row = 0 ; row < nrows ; row++ ) {
    for ( int col = 0 ; col < ncols ; col++ ) {
      bug.file.width(2);
      bug << matrix(row,col) << ",";
    }
    bug << std::endl;
  }

  
  // Apply Munkres algorithm to matrix.
  bug << "optimize allocation" << std::endl;
  Munkres m;
  m.solve(matrix);


  // Display solved matrix.
  bug << "solved matrix" << std::fflush << std::endl;
  for ( int row = 0 ; row < nrows ; row++ ) {
    for ( int col = 0 ; col < ncols ; col++ ) {
      bug.file.width(2);
      bug << matrix(row,col)<< ",";
    }
    bug << std::endl;
  }

  std::set<std::pair<uLocation,uLocation> > plan;
  for ( int row = 0 ; row != nrows ; row++ ) 
    {
      std::vector<unsigned> bestcol;
      int rowcount = 0;
      for ( int col = 0 ; col < ncols ; col++  ) 
	{
	  if ( matrix(row,col) == 0 )
	    {
	      rowcount++;
	      bestcol.push_back(col);
	    }
	}
      if ( rowcount != 1 )
	bug << "Row " << row << " has " << rowcount << " columns that have been matched." << std::endl;

      if(not bestcol.empty())
	{
	  signed select = bestcol.at(rand()%bestcol.size());
	  plan.insert(std::pair<uLocation,uLocation>(ressources.get_location(row),tasks.get_location(select)));
	}
      else
	{
	  bug << "No valid solution found for " << ressources.get_location(row) << std::endl;
	}
    }

  return plan;
}



std::ostream & dump(std::ostream & out, Map & map, const Option & opt)
{
  out << '(' << map.get_row(opt.from) << ',' << map.get_col(opt.from)  << ") -> ("
      << map.get_row(opt.to) << ',' << map.get_col(opt.to)  << ") ["
      << ("NESW-"[map.getDirection(opt.from,opt.to)]) << ']'
      << std::endl;
  return out;
}


std::ostream & dump(std::ostream & out, Map & map, const std::vector<Option> & options)
{
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    {
      out << '(' << map.get_row(it->from) << ',' << map.get_col(it->from)  << ") -> ("
	  << map.get_row(it->to) << ',' << map.get_col(it->to)  << ") ["
	  << (std::string("NESW-")[map.getDirection(it->from,it->to)]) << ']'
	  << std::endl;
    }
  
  return out;
}



std::ostream & dump(std::ostream & out, const std::vector<Option> & options)
{
  for(std::vector<Option>::const_iterator it = options.begin();
      it !=options.end(); ++it)
    {
      out << '(' << it->from  << ") -> ("
	  << it->to  << ")"
	  << std::endl;
    }
  
  return out;
}

